import cv2
import numpy as np
from ultralytics import YOLO
import os
from typing import Dict, List, Tuple
import json

class CustomYOLODetector:
    """使用自定义训练的YOLO模型进行颜色和形状检测"""
    
    def __init__(self, model_path: str, classes_file: str = None):
        """
        初始化检测器
        :param model_path: 自定义模型文件路径 (best.pt)
        :param classes_file: 类别文件路径 (可选)
        """
        try:
            # 加载自定义训练的模型
            self.model = YOLO(model_path)
            print(f"成功加载模型: {model_path}")
            
            # 从模型获取类别名称
            self.class_names = self.model.names
            print(f"模型类别: {list(self.class_names.values())}")
            
            # 如果提供了类别文件，则从文件读取
            if classes_file and os.path.exists(classes_file):
                with open(classes_file, 'r', encoding='utf-8') as f:
                    custom_classes = [line.strip() for line in f.readlines() if line.strip()]
                print(f"从文件读取的类别: {custom_classes}")
            
        except Exception as e:
            print(f"加载模型失败: {e}")
            raise
        
        # 定义颜色映射
        self.color_mapping = {
            # 长方条类别
            'bluerec': {'color': '蓝色', 'shape': '长方条'},
            'greenrec': {'color': '绿色', 'shape': '长方条'},
            'redrec': {'color': '红色', 'shape': '长方条'},
            
            # 圆柱台类别
            'bluecylinder': {'color': '蓝色', 'shape': '圆柱台'},
            'greencylinder': {'color': '绿色', 'shape': '圆柱台'},
            'redcylinder': {'color': '红色', 'shape': '圆柱台'},
        }
        
        # 验证类别映射
        self._validate_classes()
    
    def _validate_classes(self):
        """验证模型类别与颜色映射的一致性"""
        model_classes = set(self.class_names.values())
        mapped_classes = set(self.color_mapping.keys())
        
        missing_in_mapping = model_classes - mapped_classes
        extra_in_mapping = mapped_classes - model_classes
        
        if missing_in_mapping:
            print(f"警告: 以下模型类别未在颜色映射中定义: {missing_in_mapping}")
        
        if extra_in_mapping:
            print(f"警告: 以下颜色映射类别不在模型中: {extra_in_mapping}")
        
        print("类别验证完成")
    
    def get_color_and_shape(self, class_name: str) -> Tuple[str, str]:
        """
        根据类别名称获取颜色和形状
        :param class_name: 检测到的类别名称
        :return: (颜色, 形状) 元组
        """
        if class_name in self.color_mapping:
            info = self.color_mapping[class_name]
            return info['color'], info['shape']
        else:
            return '未知颜色', '未知形状'
    
    def detect_image(self, image_path: str, confidence: float = 0.5) -> Dict:
        """
        检测图像中的对象
        :param image_path: 图像路径
        :param confidence: 置信度阈值
        :return: 检测结果字典
        """
        try:
            # 读取图像
            image = cv2.imread(image_path)
            if image is None:
                return {'error': f'无法读取图像: {image_path}'}
            
            # 使用模型进行预测
            results = self.model(image, conf=confidence)
            
            # 解析检测结果
            detections = []
            
            for result in results:
                boxes = result.boxes
                if boxes is not None:
                    for i, box in enumerate(boxes):
                        # 获取边界框坐标
                        x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                        
                        # 获取置信度
                        conf = box.conf[0].cpu().numpy()
                        
                        # 获取类别ID和名称
                        class_id = int(box.cls[0].cpu().numpy())
                        class_name = self.class_names[class_id]
                        
                        # 获取颜色和形状
                        color, shape = self.get_color_and_shape(class_name)
                        
                        # 计算中心点和尺寸
                        center_x = int((x1 + x2) / 2)
                        center_y = int((y1 + y2) / 2)
                        width = int(x2 - x1)
                        height = int(y2 - y1)
                        area = width * height
                        
                        detection = {
                            'class_name': class_name,
                            'color': color,
                            'shape': shape,
                            'confidence': float(conf),
                            'bbox': [int(x1), int(y1), int(x2), int(y2)],
                            'center': [center_x, center_y],
                            'size': [width, height],
                            'area': area
                        }
                        
                        detections.append(detection)
            
            # 统计结果
            color_stats = {}
            shape_stats = {}
            
            for detection in detections:
                color = detection['color']
                shape = detection['shape']
                
                color_stats[color] = color_stats.get(color, 0) + 1
                shape_stats[shape] = shape_stats.get(shape, 0) + 1
            
            return {
                'image_path': image_path,
                'total_detections': len(detections),
                'detections': detections,
                'color_statistics': color_stats,
                'shape_statistics': shape_stats
            }
            
        except Exception as e:
            return {'error': f'检测过程中发生错误: {str(e)}'}
    
    def detect_batch(self, image_folder: str, confidence: float = 0.5) -> Dict:
        """
        批量检测文件夹中的图像
        :param image_folder: 图像文件夹路径
        :param confidence: 置信度阈值
        :return: 批量检测结果
        """
        if not os.path.exists(image_folder):
            return {'error': f'文件夹不存在: {image_folder}'}
        
        # 支持的图像格式
        image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp']
        image_files = []
        
        for file in os.listdir(image_folder):
            if any(file.lower().endswith(ext) for ext in image_extensions):
                image_files.append(os.path.join(image_folder, file))
        
        if not image_files:
            return {'error': '文件夹中没有找到支持的图像文件'}
        
        print(f"找到 {len(image_files)} 个图像文件，开始批量检测...")
        
        batch_results = {
            'total_images': len(image_files),
            'successful_detections': 0,
            'failed_detections': 0,
            'all_detections': [],
            'overall_color_stats': {},
            'overall_shape_stats': {},
            'image_results': []
        }
        
        for i, image_path in enumerate(image_files, 1):
            print(f"处理第 {i}/{len(image_files)} 个图像: {os.path.basename(image_path)}")
            
            result = self.detect_image(image_path, confidence)
            
            if 'error' in result:
                batch_results['failed_detections'] += 1
                print(f"  错误: {result['error']}")
            else:
                batch_results['successful_detections'] += 1
                batch_results['image_results'].append(result)
                batch_results['all_detections'].extend(result['detections'])
                
                # 累计统计
                for color, count in result['color_statistics'].items():
                    batch_results['overall_color_stats'][color] = \
                        batch_results['overall_color_stats'].get(color, 0) + count
                
                for shape, count in result['shape_statistics'].items():
                    batch_results['overall_shape_stats'][shape] = \
                        batch_results['overall_shape_stats'].get(shape, 0) + count
                
                print(f"  检测到 {result['total_detections']} 个对象")
        
        return batch_results
    
    def detect_video_stream(self, camera_index: int = 0, max_frames: int = 100, 
                          confidence: float = 0.5):
        """
        检测摄像头视频流
        :param camera_index: 摄像头索引
        :param max_frames: 最大处理帧数
        :param confidence: 置信度阈值
        """
        cap = cv2.VideoCapture(camera_index)
        
        if not cap.isOpened():
            print("无法打开摄像头")
            return
        
        print(f"开始视频流检测，将处理 {max_frames} 帧...")
        print("=" * 60)
        
        frame_count = 0
        total_detections = []
        
        try:
            while frame_count < max_frames:
                ret, frame = cap.read()
                if not ret:
                    break
                
                frame_count += 1
                
                # 每隔5帧检测一次
                if frame_count % 5 == 0:
                    print(f"\n第 {frame_count} 帧检测:")
                    print("-" * 40)
                    
                    # 使用模型预测
                    results = self.model(frame, conf=confidence)
                    
                    frame_detections = []
                    
                    for result in results:
                        boxes = result.boxes
                        if boxes is not None:
                            for box in boxes:
                                class_id = int(box.cls[0].cpu().numpy())
                                class_name = self.class_names[class_id]
                                conf = float(box.conf[0].cpu().numpy())
                                
                                color, shape = self.get_color_and_shape(class_name)
                                
                                detection = {
                                    'class_name': class_name,
                                    'color': color,
                                    'shape': shape,
                                    'confidence': conf
                                }
                                
                                frame_detections.append(detection)
                                total_detections.append(detection)
                    
                    if frame_detections:
                        for i, det in enumerate(frame_detections, 1):
                            print(f"  对象 {i}: {det['shape']} - {det['color']} "
                                  f"(类别: {det['class_name']}, 置信度: {det['confidence']:.2f})")
                    else:
                        print("  未检测到对象")
        
        finally:
            cap.release()
        
        # 输出总结
        if total_detections:
            print("\n" + "=" * 60)
            print("视频流检测总结")
            print("=" * 60)
            
            color_count = {}
            shape_count = {}
            
            for det in total_detections:
                color_count[det['color']] = color_count.get(det['color'], 0) + 1
                shape_count[det['shape']] = shape_count.get(det['shape'], 0) + 1
            
            print(f"总检测次数: {len(total_detections)}")
            print(f"处理帧数: {frame_count}")
            
            print("\n颜色统计:")
            for color, count in sorted(color_count.items()):
                print(f"  {color}: {count} 次")
            
            print("\n形状统计:")
            for shape, count in sorted(shape_count.items()):
                print(f"  {shape}: {count} 次")
    
    def print_detection_results(self, result: Dict):
        """
        格式化输出检测结果
        :param result: 检测结果字典
        """
        if 'error' in result:
            print(f"错误: {result['error']}")
            return
        
        print("\n" + "=" * 60)
        print(f"检测结果: {os.path.basename(result['image_path'])}")
        print("=" * 60)
        
        print(f"总检测对象数: {result['total_detections']}")
        
        if result['detections']:
            print(f"\n详细检测结果:")
            print("-" * 50)
            
            for i, detection in enumerate(result['detections'], 1):
                print(f"对象 {i}:")
                print(f"  - 形状: {detection['shape']}")
                print(f"  - 颜色: {detection['color']}")
                print(f"  - 类别: {detection['class_name']}")
                print(f"  - 置信度: {detection['confidence']:.3f}")
                print(f"  - 位置: ({detection['center'][0]}, {detection['center'][1]})")
                print(f"  - 尺寸: {detection['size'][0]} x {detection['size'][1]} 像素")
                print(f"  - 面积: {detection['area']} 平方像素")
                print()
            
            print("统计信息:")
            print("-" * 20)
            
            if result['color_statistics']:
                print("颜色分布:")
                for color, count in sorted(result['color_statistics'].items()):
                    print(f"  {color}: {count} 个")
            
            if result['shape_statistics']:
                print("\n形状分布:")
                for shape, count in sorted(result['shape_statistics'].items()):
                    print(f"  {shape}: {count} 个")
        else:
            print("\n未检测到任何对象")

def main():
    print("自定义YOLO模型检测系统")
    print("=" * 40)
    
    # 获取模型路径
    model_path = r"./best.pt"
    if not model_path:
        model_path = "best.pt"
    
    if not os.path.exists(model_path):
        print(f"模型文件不存在: {model_path}")
        return
    
    try:
        # 初始化检测器
        detector = CustomYOLODetector(model_path)
        choice = '3'
        if choice == '1':
            image_path = input("请输入图像文件路径: ").strip()
            confidence = input("请输入置信度阈值 (默认0.5): ").strip()
            confidence = float(confidence) if confidence else 0.5
            
            result = detector.detect_image(image_path, confidence)
            detector.print_detection_results(result)
        
        elif choice == '2':
            folder_path = input("请输入图像文件夹路径: ").strip()
            confidence = input("请输入置信度阈值 (默认0.5): ").strip()
            confidence = float(confidence) if confidence else 0.5
            
            batch_result = detector.detect_batch(folder_path, confidence)
            
            if 'error' in batch_result:
                print(f"错误: {batch_result['error']}")
            else:
                print(f"\n批量检测完成!")
                print(f"总图像数: {batch_result['total_images']}")
                print(f"成功检测: {batch_result['successful_detections']}")
                print(f"检测失败: {batch_result['failed_detections']}")
                print(f"总对象数: {len(batch_result['all_detections'])}")
                
                if batch_result['overall_color_stats']:
                    print("\n整体颜色统计:")
                    for color, count in sorted(batch_result['overall_color_stats'].items()):
                        print(f"  {color}: {count} 个")
                
                if batch_result['overall_shape_stats']:
                    print("\n整体形状统计:")
                    for shape, count in sorted(batch_result['overall_shape_stats'].items()):
                        print(f"  {shape}: {count} 个")
        
        elif choice == '3':
            camera_index = 0
            camera_index = int(camera_index) if camera_index.isdigit() else 0
            
            max_frames = 30
            max_frames = int(max_frames) if max_frames.isdigit() else 100
            
            confidence = 0.5
            confidence = float(confidence) if confidence else 0.5
            
            detector.detect_video_stream(camera_index, max_frames, confidence)
        
        elif choice == '4':
            print("\n模型信息:")
            print("-" * 30)
            print(f"模型路径: {model_path}")
            print(f"支持的类别: {list(detector.class_names.values())}")
            print(f"颜色映射: {json.dumps(detector.color_mapping, ensure_ascii=False, indent=2)}")
        
        else:
            print("无效选择")
    
    except Exception as e:
        print(f"程序运行错误: {e}")

if __name__ == "__main__":
    main()